from PIL import Image
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt

model_save_path = './checkpoint/mnist.ckpt'
model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])
model.load_weights(model_save_path)

# image_path = input("the path of test picture:")
image_path = '6.png'
img = Image.open(image_path)
image = plt.imread(image_path)

plt.imshow(image)

img = img.resize((28, 28), Image.ANTIALIAS)
img_arr = np.where(np.array(img.convert('L')) < 200, 255, 0) / 255.0

x_predict = img_arr[tf.newaxis, ...]
result = model.predict(x_predict)
pred = tf.argmax(result, axis=1)
print('\n')
tf.print('Predict result: ', end='')
tf.print(pred)
plt.pause(1)
plt.close()
